import datetime
import json
import traceback
import types

import numpy as np
import pandas as pd
from pandas.io.sql import has_table
from sqlalchemy import inspect, text
from sqlalchemy.orm import sessionmaker

from air_web.data_platform.mydb import MyDB
from air_web.data_platform.sql_connection import SqlAlchemyConn


class MySQLPandas(object):
    def __init__(self, **kwargs):
        self.engine = SqlAlchemyConn(**kwargs)
        self.base_model = kwargs.get("sqlalchemy_model", None)

    def _check_columns(self, columns, table_name):
        df_columns = self.read_df_by_sql(
            f"SELECT COLUMN_NAME FROM INFORMATION_SCHEMA.COLUMNS WHERE TABLE_NAME = '{table_name}';"
        )
        table_columns = (
            set([i[0] for i in df_columns.values])
            if df_columns.shape[0]
            else set()
        )
        if not isinstance(columns, list):
            return False, "columns must be list type"
        not_exist_columns = set(columns).difference(table_columns)
        if not_exist_columns:
            return False, f"columns {not_exist_columns} not exist!"
        return True, ""

    def _parse_columns_str(self, columns):
        return ",".join(["`" + i + "`" for i in columns])

    @property
    def sql_session(self):
        return SqlAlchemyConn.get_alchemy_session(self.engine)

    def read_df_by_table_name(self, table_name, columns=None, chunksize=None):
        """
        :param table_name: str, table name for matching
        :param columns: list, table columns for matching
        :param chunk_size: int, read rows one-time
        :return: pd.DataFrame
        """
        if columns:
            check_rst, err_msg = self._check_columns(columns, table_name)
            if check_rst:
                parse_columns = self._parse_columns_str(columns)
                sql = f"SELECT {parse_columns} FROM {table_name};"
            else:
                return err_msg
        else:
            sql = f"SELECT * FROM {table_name};"
        return self.read_df_by_sql(sql, chunksize=chunksize)

    def read_df_by_filed_in(
        self, field_name, field_vals, table_name, to_upper=True, **kwargs
    ):
        """
        :param field_name: str, table column for matching
        :param fiele_vals: list, values for matching
        :param table_name: str
        :return: pd.DataFrame
        """
        if to_upper:
            sql = "SELECT * from %s WHERE %s in (%s);" % (
                table_name.upper(),
                field_name.upper(),
                ",".join(field_vals),
            )
        else:
            sql = "SELECT * from %s WHERE %s in (%s);" % (
                table_name,
                field_name,
                ",".join(field_vals),
            )
        return pd.read_sql_query(sql=sql, con=self.engine, **kwargs)

    def read_df_by_sql(self, sql, params=None, index_col=None, **kwargs):
        """
        execute sql and read df
        :param sql:  sql eg:SELECT * FROM HB_INSPECTION WHERE ORG_NO = %(org_no)s
        :param params:   sql params eg: {'org_no':31401}
        :param index_col: use columns as DataFrame index
        :return:
        """
        rst = pd.read_sql_query(
            sql=sql,
            con=self.engine,
            params=params,
            index_col=index_col,
            **kwargs,
        )
        if isinstance(rst, types.GeneratorType):
            df_list = [i for i in rst]
            if df_list:
                return pd.concat(df_list)
            else:
                kwargs.pop("chunksize")
                return pd.read_sql_query(
                    sql=sql,
                    con=self.engine,
                    params=params,
                    index_col=index_col,
                    **kwargs,
                )
        return rst

    def get_class_by_tablename(self, tablename):
        """Return class reference mapped to table.

        :param tablename: String with name of table.
        :return: Class reference or None.
        """
        for c in self.base_model._decl_class_registry.values():
            if hasattr(c, "__tablename__") and c.__tablename__ == tablename:
                return c
        return None

    def check_table(self, table_name):
        res = has_table(table_name=table_name, con=self.engine)
        return res

    def all_tables(self):
        return self.query("SHOW tables")

    def query(
        self,
        sql,
        params=None,
        index_col=None,
        columns=None,
        sort_field=None,
        scroll=False,
        scroll_size=100000,
        start_index=0,
    ):
        """
        :param sql: SELECT * FROM TABLENAME;
        :param params:
        :param index_col:
        :param columns: only work * in sql
        :param sort_field: ['-a', 'b']
        :param scroll:
        :param scroll_size:
        :param start_index:
        :return:
        """
        if "*" in sql and columns:
            sql = sql.replace("*", " " + ",".join(map(str, columns)) + " ")
        # 排序
        if sort_field:
            sort_str = " ORDER BY "
            for f in sort_field:
                if f.startswith("-"):
                    sort_str += f[1:] + " DESC,"
                else:
                    sort_str += f + " ASC,"
            # 去掉最后的，
            sort_str = sort_str[:-1]
            sql += sort_str
        # 分页
        if scroll:
            sql += " LIMIT {} OFFSET {}".format(scroll_size, start_index)
        sql = sql.replace("  ", " ")
        if len(sql) < 500:
            print(sql)
        ret_df = pd.read_sql_query(sql, self.engine, params=params)
        # 保留设定索引的列, 重写 read_sql 内部逻辑
        if index_col is not None:
            ret_df.set_index(index_col, inplace=True, drop=False)
        return ret_df

    def get_single_by_id(self, table_name, data_id, index_col=None):
        ret_df = pd.read_sql_query(
            'SELECT * FROM {} WHERE id = "{}"'.format(table_name, data_id),
            self.engine,
        )
        if index_col is not None:
            ret_df.set_index(index_col, inplace=True, drop=True)
        return ret_df

    def query_by_rules(
        self,
        table_name,
        rules,
        index_col=None,
        columns=None,
        start_index=0,
        scroll=False,
        scroll_size=10000,
        sort_field=None,
        use_dumps=True,
    ):
        if not rules or not isinstance(rules, list):
            rules = []
        sql_str = "SELECT * FROM {table_name}".format(table_name=table_name)
        if rules:
            condition_list = []
            for r in rules:
                if "query" in r:
                    condition = " {field_name} {exp} {value}".format(
                        field_name=r[0],
                        exp=r[2],
                        value=json.dumps(r[3], ensure_ascii=False)
                        if use_dumps
                        else r[3],
                    )
                    if "in" in r or "in_" in r:
                        condition = condition.replace("[", "(").replace(
                            "]", ")"
                        )
                    if "in_" in r:
                        condition = condition.replace("[", "(").replace(
                            "]", ")"
                        )
                        condition = condition.replace("in_", "NOT IN")
                    if "like" in r:
                        condition = condition.replace("%", "%%")
                    condition_list.append(condition)
                if type(r) == str:
                    condition_list.append(r)
            condition_str = " AND ".join(condition_list)
            sql_str = sql_str + " WHERE" + condition_str

        ret = self.query(
            sql_str,
            index_col=index_col,
            columns=columns,
            sort_field=sort_field,
            start_index=start_index,
            scroll=scroll,
            scroll_size=scroll_size,
        )
        return ret

    def query_count_by_sql(self, sql_str, params=None):
        try:
            if not params:
                result = self.engine.execute(sql_str)
            else:
                result = self.engine.execute(sql_str, params)
            result = result.fetchall()
        except Exception as e:
            print("in query(): {}".format(e))
            print(traceback.format_exc())
            result = None
        return result[0][0] if result else 0

    def query_count_by_table_name(self, table_name):
        sql_str = "SELECT COUNT(0) FROM {table_name}".format(
            table_name=table_name
        )
        return self.query_count_by_sql(sql_str)

    def query_count_by_rules(self, table_name, rules, use_dumps=True):
        if not rules or not isinstance(rules, list):
            rules = []
        sql_str = "SELECT COUNT(0) FROM {table_name}".format(
            table_name=table_name
        )
        if rules:
            condition_list = []
            for r in rules:
                if "query" in r:
                    condition = " {field_name} {exp} {value}".format(
                        field_name=r[0],
                        exp=r[2],
                        value=json.dumps(r[3]) if use_dumps else r[3],
                    )
                    if "in" in r or "in_" in r:
                        condition = condition.replace("[", "(").replace(
                            "]", ")"
                        )
                    if "in_" in r:
                        condition = condition.replace("in_", "NOT IN")
                    if "like" in r:
                        pass
                    condition_list.append(condition)
            condition_str = " AND ".join(condition_list)
            sql_str = sql_str + " WHERE" + condition_str
        return self.query_count_by_sql(sql_str)

    def get_table_desc(self, table_name):
        return self.query("DESC {}".format(table_name))

    def insert_df(
        self,
        df,
        table_name,
        index=False,
        drop_id=True,
        chunksize=1000,
        **kwargs,
    ):
        """
        将df 数据全部添加到mysql 数据库中
        :param df: Dataframe， id should not in df.columns
        :param table_name: sql table name
        :param index:  boolean, default False Write DataFrame index as a column.
        :param drop_id: boolean, default True, 默认删除id列，当设置会False时，插入数据会有主键冲突风险。
        :param chunksize: then rows will be written in batches of this size at a time
        :return:
        """
        if df.empty:
            return
        # check_res = self.check_table(table_name)
        # if not check_res:
        #     raise Exception('sql table:{} not exists, place create it first!')
        if drop_id and "id" in df.columns:
            print("drop df columns: id")
            df = df.drop(columns=["id"], axis=1)
        df.to_sql(
            table_name,
            self.engine,
            if_exists=kwargs.pop("if_exists", "append"),
            index=index,
            chunksize=chunksize,
            **kwargs,
        )
        return True

    def insert_df_ignore(self, df, table_name, **kwargs):
        if len(df) < 1:
            return
        ignore_id = kwargs.pop("ignore_id", False)
        if not ignore_id:
            df["id"] = df.index
        df_cols = df.columns
        cols_str = self._sql_string(df_cols)
        param_symbols = self._sql_string(
            list(map(lambda col: ":" + str(col), df_cols))
        )
        text_obj = text(
            "INSERT IGNORE INTO {} {} VALUES {}".format(
                table_name, cols_str, param_symbols
            )
        )
        self.engine.execute(text_obj, self._parameterize_for_sql(df), **kwargs)

    def delete_df_by_col_value(self, df, table_name, col_name="id"):
        """
        :param df:
        :param table_name:
        :param col_name: delete by column name, default id
        :return:
        """
        values_list = list(df[col_name].values)
        if not values_list:
            return
        delete_rows = self._sql_string(df[col_name].values)
        self.engine.execute(
            text(
                "DELETE FROM {} WHERE {} IN {}".format(
                    table_name, col_name, delete_rows
                )
            )
        )

    def delete_table(self, table_name):
        self.engine.execute(text("DELETE FROM {} ".format(table_name)))

    def create_table(self, info):
        self.engine.execute(text(info))

    def truncate_table(self, table_name):
        self.engine.execute(text("TRUNCATE {} ".format(table_name)))

    def update_data_by_id(self, data_list, table_name):
        # self.check_columns(df, table_name)
        df_cols = data_list[0].keys()
        cols_str = self._sql_string(df_cols)
        param_symbols = self._sql_string(
            list(map(lambda col: ":" + str(col), df_cols))
        )
        duplicate_str = self._duplicate_str(df_cols)
        text_obj = text(
            "INSERT INTO {} {} VALUES {}  ON DUPLICATE KEY UPDATE {}".format(
                table_name, cols_str, param_symbols, duplicate_str
            )
        )
        self.engine.execute(text_obj, data_list)
        return True

    def update_df_by_id(self, df, table_name, no_update_col_list=None):
        # self.check_columns(df, table_name)
        df_cols = df.columns
        cols_str = self._sql_string(df_cols)
        param_symbols = self._sql_string(
            list(map(lambda col: ":" + str(col), df_cols))
        )
        duplicate_list = list(set(df_cols)-set(no_update_col_list)) if no_update_col_list else df_cols
        duplicate_str = self._duplicate_str(duplicate_list)
        text_obj = text(
            "INSERT INTO {} {} VALUES {}  ON DUPLICATE KEY UPDATE {}".format(
                table_name, cols_str, param_symbols, duplicate_str
            )
        )
        self.engine.execute(text_obj, self._parameterize_for_sql(df))
        return True

    def update_df_by_id_with_session(self, df, table_name):
        df_cols = df.columns
        cols_str = self._sql_string(df_cols)
        param_symbols = self._sql_string(
            list(map(lambda col: ":" + str(col), df_cols))
        )
        duplicate_str = self._duplicate_str(df_cols)
        text_obj = text(
            "INSERT INTO {} {} VALUES {}  ON DUPLICATE KEY UPDATE {}".format(
                table_name, cols_str, param_symbols, duplicate_str
            )
        )
        session = sessionmaker(bind=self.engine)()
        try:
            # 获取session并执行sql
            session.execute(text_obj, self._parameterize_for_sql(df))
            session.commit()
        except Exception as e:
            # rollback the transaction on error
            session.rollback()
            print(f"Error: {e}")
            return False
        finally:
            # close the session when you're done
            session.close()
        return True

    def insert_check_duplicate(self, df, table_name, ignore_id=False):
        if len(df) < 1:
            return
        if not ignore_id:
            df["id"] = df.index
            df_cols = df.columns
            cols_str = self._sql_string(df_cols)
            param_symbols = self._sql_string(
                list(map(lambda col: ":" + str(col), df_cols))
            )
            duplicate_str = self._duplicate_str(df_cols)
            text_obj = text(
                "INSERT INTO {} {} VALUES {}  ON DUPLICATE KEY UPDATE {}".format(
                    table_name, cols_str, param_symbols, duplicate_str
                )
            )
        else:
            df_cols = df.columns
            cols_str = self._sql_string(df_cols)
            param_symbols = self._sql_string(
                list(map(lambda col: ":" + str(col), df_cols))
            )
            text_obj = text(
                "REPLACE INTO {} {} VALUES {}".format(
                    table_name, cols_str, param_symbols
                )
            )
        self.engine.execute(text_obj, self._parameterize_for_sql(df))
        return True

    def insert_df_new(self, df, table_name):
        df = df.fillna("")
        if "id" in df.columns:
            df = df.drop(columns=["id"], axis=1)
        df_cols = df.columns
        cols_str = self._sql_string(df_cols)
        param_symbols = self._sql_string(
            list(map(lambda col: ":" + str(col), df_cols))
        )
        text_obj = text(
            "INSERT INTO {} {} VALUES {}".format(
                table_name, cols_str, param_symbols
            )
        )
        print(text_obj, self._parameterize_for_sql(df))
        self.engine.execute(text_obj, self._parameterize_for_sql(df))
        return True

    def update_df_by_col(self, df, table_name, col, op_type):
        if df.shape[0]:
            params = {
                "where_values": df[col].values.tolist(),
                "op_type": op_type,
            }
            param_symbols = "OPERATE= :op_type, VALID=0"
            sql_str = text(
                "update {} set {} where {} in :where_values and valid=1".format(
                    table_name, param_symbols, col
                )
            )
            self.engine.execute(sql_str, params)
        return True

    def update_row_by_col(self, df, table_name, col):
        for index, row in df.iterrows():
            col_val = row[col]
            row = row.drop(labels=[col])
            update_dict = row.to_dict()
            param_symbols = ", ".join(
                list(
                    map(
                        lambda col: str(col) + "= :" + str(col),
                        update_dict.keys(),
                    )
                )
            )
            sql_str = text(
                "update {} set {} where {}='{}' ".format(
                    table_name, param_symbols, col, col_val
                )
            )
            self.engine.execute(sql_str, self._format_data(update_dict))
        return True

    def update_row_by_dict(self, row_dict, table_name, where_column):
        where_value = row_dict.pop(where_column)
        if isinstance(where_value, list):
            where_str = "{} in ({})".format(
                where_column, "'" + "','".join(where_value) + "'"
            )
        else:
            where_str = "{} = {}".format(where_column, where_value)

        param_symbols = ", ".join(
            list(map(lambda col: str(col) + "= :" + str(col), row_dict.keys()))
        )
        sql_str = text(
            "update {} set {} where {} ".format(
                table_name, param_symbols, where_str
            )
        )
        self.engine.execute(sql_str, self._format_data(row_dict))

    @staticmethod
    def _format_data(row):
        rst = {}
        for k, v in row.items():
            if isinstance(v, (str, datetime.datetime)):
                rst[k] = str(v)
            else:
                rst[k] = v
        return rst

    @staticmethod
    def _duplicate_str(cols):
        cols_format = map(lambda x: "{}=VALUES({})".format(x, x), cols)
        return ",".join(cols_format)

    @staticmethod
    def _sql_string(cols):
        return "(" + ", ".join(map(str, cols)) + ")"

    def _parameterize_for_sql(self, df):
        cols = df.select_dtypes(include=np.datetime64).columns.values
        int_cols = df.select_dtypes(include=np.int64).columns.values
        float_cols = df.select_dtypes(include=np.float64).columns.values
        records = df.to_dict("records")
        for row in records:
            for col in cols:
                row[col] = str(row[col])
            for col in int_cols:
                row[col] = int(row[col])
            for col in float_cols:
                row[col] = float(row[col])
        return records

    def check_columns(self, df, table_name):
        table_model = self.get_class_by_tablename(table_name)
        inspect_obj = inspect(table_model)
        primary_key = inspect_obj.primary_key[0].name
        columns = inspect_obj.columns
        nullable_false = []
        for c in columns:
            if not c.nullable:
                nullable_false.append(c.name)
        df_columns = df.columns
        if primary_key not in df_columns:
            raise Exception(
                "primary key:{} not in dataframe".format(primary_key)
            )
        nullable_check_set = set(nullable_false) - set(df_columns)
        if nullable_check_set:
            raise Exception(
                "some not nullable:{} need in dataframe".format(
                    nullable_check_set
                )
            )
        return True


class MySQLHelper(object):
    sql_helper = None  # 原生SQL操作
    sql_pandas = None  # pandas 方式操作

    def __init__(self, init_conn=True, **kwargs):
        if init_conn:
            if not self.sql_helper:
                self.sql_helper = MyDB(
                    host=kwargs.get("MYSQL_HOST"),
                    user=kwargs.get("MYSQL_USER"),
                    pwd=kwargs.get("MYSQL_PASSWORD"),
                    port=int(kwargs.get("MYSQL_PORT")),
                    dbname=kwargs.get("MYSQL_DB_NAME"),
                    charset="utf8",
                )
            if not self.sql_pandas:
                self.sql_pandas = MySQLPandas(**kwargs)
